package cn.iocoder.yudao.framework.chatgpt.config;

import cn.hutool.core.util.StrUtil;
import com.unfbx.chatgpt.OpenAiClient;
import com.unfbx.chatgpt.OpenAiStreamClient;
import com.unfbx.chatgpt.function.KeyRandomStrategy;
import com.unfbx.chatgpt.interceptor.DefaultOpenAiAuthInterceptor;
import com.unfbx.chatgpt.interceptor.OpenAILogger;
import com.unfbx.chatgpt.interceptor.OpenAiResponseInterceptor;
import okhttp3.OkHttpClient;
import okhttp3.logging.HttpLoggingInterceptor;
import org.springframework.boot.autoconfigure.AutoConfiguration;
import org.springframework.boot.context.properties.EnableConfigurationProperties;
import org.springframework.context.annotation.Bean;

import java.net.InetSocketAddress;
import java.net.Proxy;
import java.util.concurrent.TimeUnit;

@AutoConfiguration
@EnableConfigurationProperties(ChatGPTProperties.class)
public class ChatGPTConfiguration {
    private final ChatGPTProperties chatGPTProperties;
    ChatGPTConfiguration(ChatGPTProperties webSocketProperties) {
        this.chatGPTProperties = webSocketProperties;
    }
    @Bean
    public OpenAiStreamClient openAiStreamClient(){
       return OpenAiStreamClient.builder()
                .apiKey(chatGPTProperties.getApiKeys())
                .apiHost(chatGPTProperties.getApiHost())
                .okHttpClient(customOkHttpClient()).build();
    }
    @Bean
    public OpenAiClient openAiClient(){
        return OpenAiClient.builder()
                .apiKey(chatGPTProperties.getApiKeys())
                .apiHost(chatGPTProperties.getApiHost())
                .okHttpClient(customOkHttpClient()).build();
    }

    private OkHttpClient customOkHttpClient(){
        OkHttpClient.Builder builder = new OkHttpClient.Builder();
        if(StrUtil.isNotBlank(chatGPTProperties.getProxyUrl()) && chatGPTProperties.getProxyPort() != 0){
            Proxy proxy = new Proxy(Proxy.Type.HTTP, new InetSocketAddress(chatGPTProperties.getProxyUrl(), chatGPTProperties.getProxyPort()));
            HttpLoggingInterceptor httpLoggingInterceptor = new HttpLoggingInterceptor(new OpenAILogger());
            httpLoggingInterceptor.setLevel(HttpLoggingInterceptor.Level.HEADERS);
            builder.proxy(proxy)//自定义代理
                    .addInterceptor(httpLoggingInterceptor)//自定义日志输出
                    .addInterceptor(new OpenAiResponseInterceptor());//自定义返回值拦截
        }
        DefaultOpenAiAuthInterceptor authInterceptor = new DefaultOpenAiAuthInterceptor();
        authInterceptor.setApiKey(chatGPTProperties.getApiKeys());
        authInterceptor.setKeyStrategy(new KeyRandomStrategy());
        return builder.addInterceptor(authInterceptor)
                .connectTimeout(10, TimeUnit.SECONDS)
                .writeTimeout(50, TimeUnit.SECONDS)
                .readTimeout(50, TimeUnit.SECONDS)
                .build();
    }
}
